package org.sz;

import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.dao.DataAccessException;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.ResultSetExtractor;

import javax.sql.DataSource;
import java.sql.*;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.function.Consumer;

/**
 * MySQL文档读取器
 * 使用HikariCP连接池管理数据库连接，支持高效读取数据库记录并转换为Document对象
 */
public class MySQLDocumentReader implements DocumentReader, AutoCloseable {
    
    private final MySQLResource mysqlResource;        // 数据库资源配置
    private final HikariDataSource dataSource;       // HikariCP数据源
    private final JdbcTemplate jdbcTemplate;         // Spring JDBC模板
    private final ExecutorService executorService;   // 异步任务执行器
    private int fetchSize = 500;                     // 批处理大小
    
    /**
     * 构造函数
     * @param mysqlResource 数据库资源配置
     */
    public MySQLDocumentReader(MySQLResource mysqlResource) {
        this.mysqlResource = mysqlResource;
        this.dataSource = createHikariDataSource();    // 初始化HikariCP数据源
        this.jdbcTemplate = new JdbcTemplate(dataSource); // 初始化JDBC模板
        // 初始化线程池，大小为CPU核心数的2倍
        this.executorService = Executors.newFixedThreadPool(
            Math.max(2, Runtime.getRuntime().availableProcessors() * 2));
    }
    
    /**
     * 创建HikariCP数据源
     * @return HikariDataSource实例
     */
    private HikariDataSource createHikariDataSource() {
        HikariConfig config = new HikariConfig();
        config.setJdbcUrl(mysqlResource.getJdbcUrl());         // 设置JDBC URL
        config.setUsername(mysqlResource.getUsername());       // 设置用户名
        config.setPassword(mysqlResource.getPassword());       // 设置密码
        
        // 配置连接池基本参数
        config.setMaximumPoolSize(mysqlResource.getMaximumPoolSize());
        config.setMinimumIdle(mysqlResource.getMinimumIdle());
        config.setConnectionTimeout(mysqlResource.getConnectionTimeout());
        config.setIdleTimeout(mysqlResource.getIdleTimeout());
        config.setMaxLifetime(mysqlResource.getMaxLifetime());
        
        // 添加额外连接属性（如SSL配置、时区等）
        Properties props = mysqlResource.getConnectionProperties();
        props.forEach((key, value) -> config.addDataSourceProperty((String) key, value));
        
        // MySQL性能优化配置
        config.addDataSourceProperty("cachePrepStmts", "true");     // 启用预编译语句缓存
        config.addDataSourceProperty("prepStmtCacheSize", "250");   // 预编译语句缓存大小
        config.addDataSourceProperty("prepStmtCacheSqlLimit", "2048"); // 缓存SQL最大长度
        config.addDataSourceProperty("useServerPrepStmts", "true"); // 使用服务器端预编译
        
        return new HikariDataSource(config);
    }
    
    /**
     * 获取文档列表（同步方法）
     * @return 文档列表
     */
    @Override
    public List<Document> get() {
        return jdbcTemplate.query(mysqlResource.getQuery(), new DocumentResultSetExtractor());
    }
    
    /**
     * 流式处理文档（避免大数据集内存溢出）
     * @param consumer 文档处理函数
     */
    public void streamDocuments(Consumer<Document> consumer) {
        jdbcTemplate.query(mysqlResource.getQuery(), (ResultSet rs) -> {
            while (rs.next()) {
                consumer.accept(mapResultSetToDocument(rs));
            }
            return null;
        });
    }
    
    /**
     * 异步获取文档列表
     * @return 包含文档列表的Future对象
     */
    public CompletableFuture<List<Document>> getAsync() {
        return CompletableFuture.supplyAsync(this::get, executorService);
    }
    
    /**
     * 设置批处理大小（优化大数据集读取）
     * @param fetchSize 批处理大小
     */
    public void setFetchSize(int fetchSize) {
        this.fetchSize = fetchSize;
        // 配置HikariCP的fetchSize
        mysqlResource.getConnectionProperties().setProperty("fetchSize", String.valueOf(fetchSize));
    }
    
    /**
     * 将ResultSet映射为Document对象
     * @param rs 结果集
     * @return Document对象
     * @throws SQLException SQL执行异常
     */
    private Document mapResultSetToDocument(ResultSet rs) throws SQLException {
        Map<String, Object> rowData = new HashMap<>();
        ResultSetMetaData metaData = rs.getMetaData();
        int columnCount = metaData.getColumnCount();
        
        // 提取行数据
        for (int i = 1; i <= columnCount; i++) {
            String columnName = metaData.getColumnName(i);
            rowData.put(columnName, rs.getObject(i));
        }
        
        // 构建文档内容和元数据
        String content = buildContent(rowData);
        Map<String, Object> metadata = buildMetadata(rowData);
        
        return new Document(content, metadata);
    }
    
    /**
     * 构建文档内容
     * @param rowData 行数据
     * @return 文档内容字符串
     */
    private String buildContent(Map<String, Object> rowData) {
        StringBuilder contentBuilder = new StringBuilder();
        List<String> contentColumns = mysqlResource.getTextColumns();
        
        // 根据配置决定使用哪些列构建内容
        if (contentColumns == null || contentColumns.isEmpty()) {
            // 无指定内容列时使用所有列
            for (Map.Entry<String, Object> entry : rowData.entrySet()) {
                appendColumnContent(contentBuilder, entry.getKey(), entry.getValue());
            }
        } else {
            // 使用指定的内容列
            for (String column : contentColumns) {
                if (rowData.containsKey(column)) {
                    appendColumnContent(contentBuilder, column, rowData.get(column));
                }
            }
        }
        return contentBuilder.toString().trim();
    }
    
    /**
     * 构建文档元数据
     * @param rowData 行数据
     * @return 元数据Map
     */
    private Map<String, Object> buildMetadata(Map<String, Object> rowData) {
        Map<String, Object> metadata = new HashMap<>();
        metadata.put(MySQLResource.SOURCE, mysqlResource.getJdbcUrl()); // 数据源标识
        
        // 添加指定的元数据列
        List<String> metadataColumns = mysqlResource.getMetadataColumns();
        if (metadataColumns != null) {
            for (String column : metadataColumns) {
                if (rowData.containsKey(column)) {
                    metadata.put(column, rowData.get(column));
                }
            }
        }
        
        // 添加额外元数据（如时间戳）
        metadata.put("timestamp", System.currentTimeMillis());
        metadata.put("fetchSize", fetchSize);
        return metadata;
    }
    
    /**
     * 追加列内容到字符串构建器
     * @param builder 字符串构建器
     * @param column 列名
     * @param value 列值
     */
    private void appendColumnContent(StringBuilder builder, String column, Object value) {
        if (value != null) {
            builder.append(column).append(": ").append(value).append("\n");
        }
    }
    
    /**
     * ResultSetExtractor实现类，用于将结果集转换为文档列表
     */
    private class DocumentResultSetExtractor implements ResultSetExtractor<List<Document>> {
        @Override
        public List<Document> extractData(ResultSet rs) throws SQLException, DataAccessException {
            List<Document> documents = new ArrayList<>();
            while (rs.next()) {
                documents.add(mapResultSetToDocument(rs));
            }
            return documents;
        }
    }
    
    /**
     * 关闭资源（实现AutoCloseable接口）
     */
    @Override
    public void close() {
        // 关闭线程池
        if (executorService != null && !executorService.isShutdown()) {
            executorService.shutdown();
        }
        // 关闭数据源
        if (dataSource != null && !dataSource.isClosed()) {
            dataSource.close();
        }
    }
}